package com.dtdream.emr;

import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.exec.UDFArgumentLengthException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.StringObjectInspector;

import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.Map;
import java.util.TreeMap;
import java.util.regex.Pattern;



public final class Ip2region extends GenericUDF {
    private transient ObjectInspector[] argumentOIs;
    private StringObjectInspector soi0;
    private StringObjectInspector soi1;
    private static TreeMap<Long, String> ipmap = null;

    @Override
    public ObjectInspector initialize(ObjectInspector[] args)
            throws UDFArgumentException {
        if (args.length > 2) {
            throw new UDFArgumentLengthException(
                    "The operator 'Ip2region' needs 2 arguments.");
        }

        argumentOIs = args;
        soi0 = (StringObjectInspector) argumentOIs[0];
        soi1 = (StringObjectInspector) argumentOIs[1];

        return PrimitiveObjectInspectorFactory.javaStringObjectInspector;
    }

    @Override
    public String getDisplayString(String[] children) {
        return "Ip2region usage: select ip2region(column) from table where ...";
    }

    private static synchronized void init_ipmap() {
        if (ipmap == null) {
            TreeMap<Long, String> thisIpMap = new TreeMap<Long, String>();

            String fileName = "ipdata_code_with_maxmind.txt";
            InputStream is = Ip2region.class.getClassLoader().getResourceAsStream(fileName);
            try {
                BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(is, "UTF-8"));

                String line;
                while ((line = bufferedReader.readLine()) != null) {
                    String[] as = line.split(",");
                    long ip_start = Long.parseLong(as[0]);
                    thisIpMap.put(ip_start, line);
                }
                bufferedReader.close();
            } catch (Exception e) {
                e.printStackTrace();
            }
            ipmap = thisIpMap;
        }

        return ;
    }

    @Override
    public Object evaluate(DeferredObject[] args) throws HiveException {

        if (null == ipmap ) {
            init_ipmap();
        }

        String ip = soi0.getPrimitiveJavaObject(args[0].get());
        String level = soi1.getPrimitiveJavaObject(args[1].get());
        String result;

        if (!Pattern.matches("(\\d{1,3}\\.){3}\\d{1,3}", ip)) {
            return null;
        }

        result = toRegion(ip, level);

        return result;
    }

    private static long ip2long(String ip_str) {
        String[] ip_array = ip_str.split("\\.");

        return (Long.parseLong(ip_array[0]) << 24) +
                (Long.parseLong(ip_array[1]) << 16) +
                (Long.parseLong(ip_array[2]) << 8) +
                Long.parseLong(ip_array[3]);
    }

    private static String long2ip(long ip) {
        StringBuilder sb_start = new StringBuilder();
        sb_start.append((ip >> 24) & 0xFF).append(".");
        sb_start.append((ip >> 16) & 0xFF).append(".");
        sb_start.append((ip >> 8) & 0xFF).append(".");
        sb_start.append(ip & 0xFF);

        return sb_start.toString();
    }

    private String toRegion(String ip, String level) {
        String result = null;

        if (null == ip) {
            return "NULL";
        }

        long ip_long = ip2long(ip);
        Map.Entry me = ipmap.floorEntry(ip_long);
        if (me == null) {
            System.out.printf("Failed to find your ip: %s %d.\r\n", ip, ip_long);
            return null;
        }

        String[] regions = ((String) me.getValue()).split(",");

        switch (level) {
            case "country":
                result = regions[2];
                break;
            case "province":
                result = regions[3];
                break;
            case "city":
                result = regions[4];
                break;
            case "district":
                result = regions[5];
                break;
            case "school":
                //??
                break;
            default:
        }

        if ("-1".equals(result)) {
            result = null;
        }

        return result;
    }
}
